#!/usr/bin/env python

"""\
A utility to edit hadoop configuration files.

Usage::

  $ edit_conf conf/yarn-site.xml tmp.xml \
       yarn.nodemanager.resource.cpu-vcores 2 \
       yarn.nodemanager.resource.memory-mb 1024
"""

from lxml import etree as ET
import sys


def doc_to_dict(doc):
    props = {}
    root = doc.getroot()
    for p in root.findall('property'):
        props[p.find('name').text] = p.find('value').text
    return props


def dict_to_doc(props):
    doc = ET.ElementTree(ET.fromstring('<configuration/>'))
    root = doc.getroot()
    pi = ET.ProcessingInstruction(
        'xml-stylesheet',
        'type="text/xsl" href="configuration.xsl"')
    root.addprevious(pi)
    for k in props:
        p = ET.SubElement(root, "property")
        name = ET.SubElement(p, "name")
        val = ET.SubElement(p, "value")
        name.text, val.text = k, props[k]
    return doc


def main(argv):
    assert len(argv) >= 2 and not (len(argv) & 0x01)
    conf_input = argv[0]
    conf_output = argv[1]
    doc = ET.parse(conf_input)
    props = doc_to_dict(doc)
    ai = iter(argv[2:])
    for k, v in zip(ai, ai):
        props[k] = v
    ndoc = dict_to_doc(props)
    with open(conf_output, 'wb') as f:
        f.write(ET.tostring(
            ndoc,
            encoding="utf-8",
            xml_declaration=True,
            pretty_print=True
        ))


if __name__ == "__main__":
    main(sys.argv[1:])
